from mpi4py import MPI
from v3dsolver.tree_manager.distribution_manager import get_key

class LoadBalancer:
    """
    Generic class for performing load balancing by redistributing particles 
    among different processors.
    """
    def __init__(self,this_elems):
        """
        Initialization of Load Balancer class
        
        this_elems: sequence containing elems of current processor
        """
        self.comm = MPI.COMM_WORLD
        self.this_elems = this_elems
    
    def get_nshift(self,src_cost,dest_cost,ns,nd):
        """
        Returns the number of particles which should be shifted from src to 
        dest. Decides this based on src_cost and dest_cost, and also based on 
        already existing number of particles in src and dest: ns and nd.
        
        Formula for calculating nshift is arrived at by calculating cost per 
        elem for each processor and ensuring that number of particles in each 
        proc is such that the total cost in both src and dest are equal.
        """
        return int((src_cost - dest_cost)*ns*nd/(ns*dest_cost+nd*src_cost))
    
    def shift_elems(self,src,dest,nshift):
        """ 
        Shifts nshift number of particles from processor of rank src to 
        processor of rank dest. Shifts this way only between processors which
        are neighbours to each other in SFC order.
        """
        if dest==src+1:
            # dest is right neighbour of src in SFC order, shift rightmost elems
            if self.comm.rank==src:
                self.comm.send(self.this_elems[-nshift:],dest,tag=9)
                self.this_elems = self.this_elems[:len(self.this_elems)-nshift]
            elif self.comm.rank==dest:
                self.this_elems += self.comm.recv(source=src,tag=9)
        
        if dest==src-1:
            # dest is left neighbour of src in SFC order, shift leftmost elems
            if self.comm.rank==src:
                self.comm.send(self.this_elems[:nshift],dest,tag=10)
                self.this_elems = self.this_elems[nshift:]
            elif self.comm.rank==dest:
                self.this_elems += self.comm.recv(source=src,tag=10)
    
    def balance_load(self,cost):
        """
        Balance load among processors by particle shifting based on given cost.
        """
        allcosts = self.comm.allgather(cost)
        all_nelems = self.comm.allgather(len(self.this_elems))
        print allcosts,all_nelems
        for i in range(len(allcosts)):
            dest = i
            if not i==self.comm.size-1: 
                if allcosts[i]>allcosts[i+1]:
                    dest = i+1
            if not i==0:
                if allcosts[i]>allcosts[i-1]:
                    dest = i-1
            if not dest==i:
                src_cost,dest_cost = allcosts[i],allcosts[dest]
                ns,nd = all_nelems[i],all_nelems[dest]
                nshift = self.get_nshift(src_cost,dest_cost,ns,nd)
                print nshift, len(self.this_elems)
                self.shift_elems(src=i,dest=dest,nshift=nshift)